"""
Author: Morphlng
Date: 2024-03-18 14:13:36
LastEditTime: 2025-04-21 09:11:54
LastEditors: Morphlng
Description: Collect and buffer data from adapter
FilePath: /DrivingGym/src/driving_gym/data/data_provider.py
"""

from __future__ import annotations

from collections import defaultdict, deque
from itertools import islice
from typing import NamedTuple

import numpy as np
from typing_extensions import TypedDict

from driving_gym.simulation.adapter_interface import AdapterInterface
from driving_gym.simulation.common.geom import BoundingBox, Transform, Vector3D


class Snapshot(NamedTuple):
    class ActorData(TypedDict):
        transform: Transform
        velocity: Vector3D
        acceleration: Vector3D
        bounding_box: BoundingBox
        sensors: dict[str, np.ndarray]

    frame: int
    """Frame number."""

    data: dict[str, ActorData]
    """Data of the frame. E.g. {"vehicle1": {"transform": ..., "velocity": ...}, ...}"""


class DataProvider:
    def __init__(self, adapter: AdapterInterface, buffer_size: int = 2000):
        self.adapter = adapter
        self.buffer = deque(maxlen=buffer_size)

    @property
    def last_frame(self) -> Snapshot:
        return self.buffer[0] if self.buffer else Snapshot(-1, {})

    def get_frames(self, n: int, step: int = 1, sort: bool = True) -> list[Snapshot]:
        """Returns the last n frames from the buffer.

        Args:
            n (int): The number of frames to return.
            step (int, optional): The step size. Defaults to 1.
            sort (bool, optional): Sort the frames by frame number. Defaults to True.

        Returns:
            list[tuple[int, dict]]: The last n frames from the buffer.
            E.g. [(100, {}), (102, {}), ...]

        Note:
            It is possible to get less than n frames if the buffer is not enough.
        """
        frames = list(islice(self.buffer, 0, n * step, step))
        if sort:
            frames.reverse()
        return frames

    def reset(self):
        self.buffer.clear()

    def collect_frame_data(self, frame: int):
        """Collects data from the adapter and buffers it.

        Args:
            frame (int): The frame returned by `adapter.tick()`.
        """
        if frame <= self.last_frame[0]:
            return

        frame_data = defaultdict(dict)
        actors = self.adapter.get_actors()
        sensors = self.adapter.get_sensors()

        for actor in actors:
            data = Snapshot.ActorData(
                {
                    "transform": self.adapter.get_transform(actor),
                    "velocity": self.adapter.get_velocity(actor),
                    "acceleration": self.adapter.get_acceleration(actor),
                    "bounding_box": self.adapter.get_bounding_box(actor),
                    "sensors": {},
                }
            )

            for sensor in filter(lambda name: name.startswith(actor), sensors):
                sensor_data = self.adapter.get_sensor_data(sensor)
                data["sensors"][sensor[len(actor) + 1 :]] = sensor_data

            frame_data[actor].update(data)

        for sensor in filter(lambda name: name.startswith("global"), sensors):
            sensor_data = self.adapter.get_sensor_data(sensor)
            frame_data["global"]["sensors"][sensor[7:]] = sensor_data

        self.buffer.appendleft(Snapshot(frame, frame_data))
